//===-- Application to analyze benchmark JSON files -----------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "automemcpy/ResultAnalyzer.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/StringSet.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/JSON.h"
#include "llvm/Support/MemoryBuffer.h"

namespace llvm {

// User can specify one or more json filenames to process on the command line.
static cl::list<std::string> InputFilenames(cl::Positional, cl::OneOrMore,
                                            cl::desc("<input json files>"));

// User can filter the distributions to be taken into account.
static cl::list<std::string>
    KeepOnlyDistributions("keep-only-distributions",
                          cl::desc("<comma separated list of distribution "
                                   "names, keeps all if unspecified>"));

namespace automemcpy {

// This is defined in the autogenerated 'Implementations.cpp' file.
extern ArrayRef<NamedFunctionDescriptor> getFunctionDescriptors();

// Iterates over all functions and fills a map of function name to function
// descriptor pointers.
static StringMap<const FunctionDescriptor *> createFunctionDescriptorMap() {
  StringMap<const FunctionDescriptor *> Descriptors;
  for (const NamedFunctionDescriptor &FD : getFunctionDescriptors())
    Descriptors.insert_or_assign(FD.Name, &FD.Desc);
  return Descriptors;
}

// Retrieves the function descriptor for a particular function name.
static const FunctionDescriptor &getFunctionDescriptor(StringRef FunctionName) {
  static StringMap<const FunctionDescriptor *> Descriptors =
      createFunctionDescriptorMap();
  const auto *FD = Descriptors.lookup(FunctionName);
  if (!FD)
    report_fatal_error(
        Twine("No FunctionDescriptor for ").concat(FunctionName));
  return *FD;
}

// Functions and distributions names are stored quite a few times so it's more
// efficient to internalize these strings and refer to them through 'StringRef'.
static StringRef getInternalizedString(StringRef VolatileStr) {
  static llvm::StringSet StringCache;
  return StringCache.insert(VolatileStr).first->getKey();
}

// Helper function for the LLVM JSON API.
bool fromJSON(const json::Value &V, Sample &Out, json::Path P) {
  std::string Label;
  std::string RunType;
  json::ObjectMapper O(V, P);
  if (O && O.map("bytes_per_second", Out.BytesPerSecond) &&
      O.map("run_type", RunType) && O.map("label", Label)) {
    const auto LabelPair = StringRef(Label).split(',');
    Out.Id.Function.Name = getInternalizedString(LabelPair.first);
    Out.Id.Function.Type = getFunctionDescriptor(LabelPair.first).Type;
    Out.Id.Distribution.Name = getInternalizedString(LabelPair.second);
    Out.Type = StringSwitch<SampleType>(RunType)
                   .Case("aggregate", SampleType::AGGREGATE)
                   .Case("iteration", SampleType::ITERATION);
    return true;
  }
  return false;
}

// An object to represent the content of the JSON file.
// This is easier to parse/serialize JSON when the structures of the json file
// maps the structure of the object.
struct JsonFile {
  std::vector<Sample> Samples;
};

// Helper function for the LLVM JSON API.
bool fromJSON(const json::Value &V, JsonFile &JF, json::Path P) {
  json::ObjectMapper O(V, P);
  return O && O.map("benchmarks", JF.Samples);
}

// Global object to ease error reporting, it consumes errors and crash the
// application with a meaningful message.
static ExitOnError ExitOnErr;

// Main JSON parsing method. Reads the content of the file pointed to by
// 'Filename' and returns a JsonFile object.
JsonFile parseJsonResultFile(StringRef Filename) {
  auto Buf = ExitOnErr(errorOrToExpected(
      MemoryBuffer::getFile(Filename, /*bool IsText=*/true,
                            /*RequiresNullTerminator=*/false)));
  auto JsonValue = ExitOnErr(json::parse(Buf->getBuffer()));
  json::Path::Root Root;
  JsonFile JF;
  if (!fromJSON(JsonValue, JF, Root))
    ExitOnErr(Root.getError());
  return JF;
}

// Serializes the 'GradeHisto' to the provided 'Stream'.
static void Serialize(raw_ostream &Stream, const GradeHistogram &GH) {
  static constexpr std::array<StringRef, 9> kCharacters = {
      " ", "▁", "▂", "▃", "▄", "▅", "▆", "▇", "█"};

  const size_t Max = *std::max_element(GH.begin(), GH.end());
  for (size_t I = 0; I < GH.size(); ++I) {
    size_t Index = (float(GH[I]) / Max) * (kCharacters.size() - 1);
    Stream << kCharacters.at(Index);
  }
}

int Main(int argc, char **argv) {
  ExitOnErr.setBanner("Automemcpy Json Results Analyzer stopped with error: ");
  cl::ParseCommandLineOptions(argc, argv, "Automemcpy Json Results Analyzer\n");

  // Reads all samples stored in the input JSON files.
  std::vector<Sample> Samples;
  for (const auto &Filename : InputFilenames) {
    auto Result = parseJsonResultFile(Filename);
    llvm::append_range(Samples, Result.Samples);
  }

  if (!KeepOnlyDistributions.empty()) {
    llvm::StringSet ValidDistributions;
    ValidDistributions.insert(KeepOnlyDistributions.begin(),
                              KeepOnlyDistributions.end());
    llvm::erase_if(Samples, [&ValidDistributions](const Sample &S) {
      return !ValidDistributions.contains(S.Id.Distribution.Name);
    });
  }

  // Extracts median of throughputs.
  std::vector<FunctionData> Functions = getThroughputs(Samples);
  fillScores(Functions);
  castVotes(Functions);

  // Present data by function type, Grade and Geomean of scores.
  std::sort(Functions.begin(), Functions.end(),
            [](const FunctionData &A, const FunctionData &B) {
              const auto Less = [](const FunctionData &FD) {
                return std::make_tuple(FD.Id.Type, FD.FinalGrade,
                                       -FD.ScoresGeoMean);
              };
              return Less(A) < Less(B);
            });

  // Print result.
  for (const FunctionData &Function : Functions) {
    outs() << formatv("{0,-10}", Grade::getString(Function.FinalGrade));
    outs() << " |";
    Serialize(outs(), Function.GradeHisto);
    outs() << "| ";
    outs().resetColor();
    outs() << formatv("{0,+25}", Function.Id.Name);
    outs() << "\n";
  }

  return EXIT_SUCCESS;
}

} // namespace automemcpy
} // namespace llvm

int main(int argc, char **argv) { return llvm::automemcpy::Main(argc, argv); }
